{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from options.test_options import TestOptions\n",
    "from data import create_dataset\n",
    "from models import create_model\n",
    "from util.visualizer import save_images\n",
    "from util import html\n",
    "import util.util as util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------- Options ---------------\n",
      "                 CUT_mode: CUT                           \n",
      "               batch_size: 1                             \n",
      "          checkpoints_dir: ./checkpoints                 \n",
      "                crop_size: 256                           \n",
      "                 dataroot: ./datasets/middle_cityRH      \t[default: placeholder]\n",
      "             dataset_mode: unaligned                     \n",
      "                direction: AtoB                          \n",
      "          display_winsize: 256                           \n",
      "               easy_label: experiment_name               \n",
      "                    epoch: latest                        \n",
      "                     eval: False                         \n",
      "        flip_equivariance: False                         \n",
      "                  gpu_ids: 0                             \n",
      "                init_gain: 0.02                          \n",
      "                init_type: xavier                        \n",
      "                 input_nc: 3                             \n",
      "                  isTrain: False                         \t[default: None]\n",
      "            lambda_DisNCE: 1.0                           \n",
      "               lambda_GAN: 1.0                           \n",
      "               lambda_MSE: 1.0                           \n",
      "               lambda_NCE: 1.0                           \n",
      "                load_size: 256                           \n",
      "         max_dataset_size: inf                           \n",
      "                    model: DualExtractor                 \t[default: cut]\n",
      "               n_layers_D: 3                             \n",
      "                     name: middle_dual_ps2               \t[default: experiment_name]\n",
      "                    nce_T: 0.07                          \n",
      "                  nce_idt: True                          \n",
      "nce_includes_all_negatives_from_minibatch: False                         \n",
      "               nce_layers: 0,4,8,12,16                   \n",
      "                      ndf: 64                            \n",
      "                     netD: basic                         \n",
      "                     netF: mlp_patch_square              \n",
      "                  netF_nc: 256                           \n",
      "                     netG: resnet_9blocks                \n",
      "                      ngf: 64                            \n",
      "             no_antialias: False                         \n",
      "          no_antialias_up: False                         \n",
      "               no_dropout: True                          \n",
      "                  no_flip: False                         \n",
      "                    normD: instance                      \n",
      "                    normG: instance                      \n",
      "              num_patches: 128                           \n",
      "                 num_test: 50                            \n",
      "              num_threads: 4                             \n",
      "                output_nc: 3                             \n",
      "                    phase: test                          \n",
      "                pool_size: 0                             \n",
      "               preprocess: resize_and_crop               \n",
      "         random_scale_max: 3.0                           \n",
      "              results_dir: ./results/                    \n",
      "        sample_patch_size: 4                             \n",
      "           serial_batches: False                         \n",
      "stylegan2_G_num_downsampling: 1                             \n",
      "                   suffix:                               \n",
      "                  verbose: False                         \n",
      "----------------- End -------------------\n",
      "dataset [UnalignedDataset] was created\n"
     ]
    }
   ],
   "source": [
    "TestOption = TestOptions('--dataroot ./datasets/middle_cityRH --model DualExtractor --name middle_dual_ps2')\n",
    "opt = TestOption.parse()  # get test options\n",
    "opt.num_threads = 0   # test code only supports num_threads = 1\n",
    "opt.batch_size = 1    # test code only supports batch_size = 1\n",
    "opt.serial_batches = True  # disable data shuffling; comment this line if results on randomly chosen images are needed.\n",
    "opt.no_flip = True    # no flip; comment this line if results on flipped images are needed.\n",
    "opt.display_id = -1   # no visdom display; the test code saves the results to a HTML file.\n",
    "dataset = create_dataset(opt)  # create a dataset given opt.dataset_mode and other options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset [UnalignedDataset] was created\n",
      "model [DUALEXTRACTORModel] was created\n"
     ]
    }
   ],
   "source": [
    "train_dataset = create_dataset(util.copyconf(opt, phase=\"train\"))\n",
    "model = create_model(opt)      # create a model given opt.model and other options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading the model from ./checkpoints/middle_dual_ps2/latest_net_Rain.pth\n",
      "loading the model from ./checkpoints/middle_dual_ps2/latest_net_Back.pth\n",
      "---------- Networks initialized -------------\n",
      "[Network Rain] Total number of parameters : 11.378 M\n",
      "[Network Back] Total number of parameters : 11.378 M\n",
      "-----------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "for i, data in enumerate(dataset):\n",
    "    model.data_dependent_initialize(data)\n",
    "    model.setup(opt)               # regular setup: load and print networks; create schedulers\n",
    "    model.parallelize()\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n"
     ]
    }
   ],
   "source": [
    "for i, data in enumerate(dataset):    \n",
    "    model.set_input(data)  \n",
    "    feature = model.netRain(model.pred_R, model.nce_layers, encode_only=True)\n",
    "    print(len(feature))\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 262, 262])\n",
      "torch.Size([1, 128, 256, 256])\n",
      "torch.Size([1, 256, 128, 128])\n",
      "torch.Size([1, 256, 64, 64])\n",
      "torch.Size([1, 256, 64, 64])\n"
     ]
    }
   ],
   "source": [
    "for f in feature:\n",
    "    print(f.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, data in enumerate(dataset):\n",
    "    if i == 0:\n",
    "        model.data_dependent_initialize(data)\n",
    "        model.setup(opt)               # regular setup: load and print networks; create schedulers\n",
    "        model.parallelize()\n",
    "        if opt.eval:\n",
    "            model.eval()\n",
    "    if i >= opt.num_test:  # only apply our model to opt.num_test images.\n",
    "        break\n",
    "    model.set_input(data)  # unpack data from data loader\n",
    "    model.test()           # run inference\n",
    "    visuals = model.get_current_visuals()  # get image results\n",
    "    img_path = model.get_image_paths()     # get image paths\n",
    "    if i % 5 == 0:  # save images to an HTML file\n",
    "        print('processing (%04d)-th image... %s' % (i, img_path))\n",
    "    save_images(webpage, visuals, img_path, width=opt.display_winsize)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
